import re
from constant import *


class RawData:
    def __init__(self, data_dir):
        self.n_lines_in_file = 0   # 整个文件中有多少行，其中每一行是 sentence, labels
        self.n_multi_aspect_span = 0  # 统计整个数据集中有多少个样本有多个 aspect
        self.n_label = 0  # 整个文件中，有多少个 (aspect, opinion, sentiment)
        self.n_multi_aspect_label = 0  # abc

        self.data = []
        # data = [{ sentence: str, triplets: [(aspect, opinions, sentiment), (aspect, opinions, sentiment), ()]
        # }, {}, {}]
        self.multi_aspect_id = []
        self.data_dir = data_dir

        # 构造 self.data
        data_lines = open(data_dir)
        for i, text in enumerate(data_lines):
            # I charge it at night and skip taking the cord with me because of the good battery life .####[([16, 17], [15], 'POS')]
            self.n_lines_in_file += 1
            sentence = re.search(r'.+(?=####)', text)[0]
            labels = eval(re.search(r'(?<=####).+', text)[0])
            # print("sentence, labels", sentence, labels)
            text_and_label_dict = {
                'sentence': sentence,
                'triplets': []
            }
            aspect_span_list, opinion_span_list, sentiment_list = [], [], []

            for aspect_spans, opinion_spans, sentiment_str in labels:
                aspect_span_left, aspect_span_right = aspect_spans[0], aspect_spans[-1]
                opinion_span_left, opinion_span_right = opinion_spans[0], opinion_spans[-1]
                if [aspect_span_left, aspect_span_right] in aspect_span_list:
                    index = aspect_span_list.index([aspect_span_left, aspect_span_right])
                    opinion_span_list[index].append([opinion_span_left, opinion_span_right])
                    if sentiment_str != sentiment_list[index]:
                        print('inconsistent sentiment')  # 此时，对于一个 aspect，其 opinion 不相同，且情感不同
                        pass
                    pass
                else:
                    aspect_span_list.append([aspect_span_left, aspect_span_right])
                    opinion_span_list.append([[opinion_span_left, opinion_span_right]])
                    sentiment_list.append(sentiment_str)
                    pass
                self.n_label += 1
                pass

            if len(aspect_span_list) > 1:  # multiple aspect
                self.multi_aspect_id.append(i)
                self.n_multi_aspect_span += 1
                self.n_multi_aspect_label += len(labels)
                pass

            for j in range(len(aspect_span_list)):
                new_triplet = (aspect_span_list[j], opinion_span_list[j], sentiment_list[j])
                text_and_label_dict['triplets'].append(new_triplet)
                pass
            self.data.append(text_and_label_dict)
            pass

        self.len = len(self.data)

        # 构造 self.all_triplets, self.all_multi_triplets, self.all_single_triplets
        self.all_triplets, self.all_multi_triplets, self.all_single_triplets = self.get_all_triplets_set()
        # 上边的这三个都是 i-al-ar-ol-or-s 格式的 str list

        # 下面三个变量好像并没有什么用
        sample_t, sample_2t, sample_2et = 0, 0, 0
        for d in self.data:
            aspect_num = len(d['triplets'])
            sample_t += 1
            sample_2t += 2 * aspect_num
            sample_2et += 2 ** aspect_num
            pass

        print(f"RawData in {data_dir} __init__() finished!")
        pass

    def __iter__(self):
        for i in range(self.len):
            yield {
                'sentence': self.data[i]['sentence'],
                'triplets': self.data[i]['triplets']
            }
            pass
        pass

    #  i-al-ar-ol-or-s set
    def get_all_triplets_set(self):
        r = ['', '', '', '', '', '']
        golden_set, multi_set, single_set = set(), set(), set()
        data_lines = open(self.data_dir)
        for i, text in enumerate(data_lines):
            sentence = re.search(r'.+(?=####)', text)[0]
            labels = eval(re.search(r'(?<=####).+', text)[0])
            r[0] = str(i)
            for j in labels:
                r[1], r[2] = str(j[0][0]), str(j[0][-1])
                r[3], r[4] = str(j[1][0]), str(j[1][-1])
                r[5] = str(sentiment2id[j[2]])
                golden_set.add('-'.join(r))
                if i in self.multi_aspect_id:
                    multi_set.add('-'.join(r))
                    pass
                else:
                    single_set.add('-'.join(r))
                    pass
                pass
            pass
        return golden_set, multi_set, single_set

    #  i-al-ar set
    def get_all_aspect_set(self):
        triplets_golden_set = self.all_triplets
        golden_set = set()
        for i in triplets_golden_set:
            t = i.split('-')[:3]
            golden_set.add('-'.join(t))
            pass
        return golden_set

    # i-ol-or set
    def get_all_opinion_set(self):
        triplets_golden_set = self.all_triplets
        golden_set = set()
        for i in triplets_golden_set:
            t = i.split('-')[0:1]
            t.extend(i.split('-')[3:5])
            golden_set.add('-'.join(t))
            pass
        return golden_set

    # i-al-ar-ol-or set
    def get_all_pair_set(self):
        triplets_golden_set = self.all_triplets
        golden_set = set()
        for i in triplets_golden_set:
            t = i.split('-')[0:-1]
            golden_set.add('-'.join(t))
            pass
        return golden_set

    # i-al-ar-s set
    # 其中 i 代表序号，al=aspect_left，ar=aspect_right，s=sentiment
    def get_all_as_set(self):
        triplets_golden_set = self.all_triplets
        golden_set = set()
        for triplet_str in triplets_golden_set:
            triplet_str_split_res = triplet_str.split('-')
            t = triplet_str_split_res[0:3]
            t.append(triplet_str_split_res[-1])
            golden_set.add('-'.join(t))
            pass
        return golden_set

    # dict {i':'al-ar-ol-or-s'}
    def get_triplets_dict(self):
        triplets_golden_list = list(self.all_triplets)
        r = {}
        for t in triplets_golden_list:
            key = t.split('-')[0]
            value = '-'.join(t.split('-')[1:])
            if key not in r.keys():
                r[key] = [value]
                pass
            else:
                r[key].append(value)
                pass
            pass
        return r

    def __len__(self):
        return self.len

    pass


def generate_union_data_from_raw_data(original_data: RawData):
    return [
        original_data.get_all_aspect_set(),
        original_data.get_all_opinion_set(),
        (
            original_data.all_triplets,
            original_data.all_multi_triplets,
            original_data.all_single_triplets,
            original_data.multi_aspect_id
        ),
        original_data.get_all_pair_set(),
        original_data.get_all_as_set(),
        original_data.multi_aspect_id,
        original_data.get_triplets_dict()
    ]


if __name__ == '__main__':
    raw_data = RawData("../../Dataset/ASTE-Data-V2/14lap/train_triplets.txt")
    print(len(raw_data))
    pass
